GatherD

按照给定的维度 dim 和索引张量 index,从输入张量中按元素位置 抽取数据,生成新的输出张量。

该算子等价于在指定维度上执行逐元素 Gather 操作,其输出形状与 index 张量形状一致。

\[\text{output}[i_0, \dots, i_n] = \text{input}_x[i_0, \dots, i_{dim-1}, \text{index}[i_0, \dots, i_n], i_{dim+1}, \dots, i_n]\]
输入:
  • input_x - 输入张量的数据地址。 数据类型需与所调用的 GatherD 接口类型一致。

  • dim - 指定进行 Gather 操作的维度索引,取值范围为 [0, input_shape_size)

  • index - 索引张量的数据地址,类型为 int*, 用于指定在 dim 维度上的取值位置。

  • input_shape - 输入张量各维度大小数组地址。

  • input_shape_size - 输入张量的维度数量。

  • index_shape - 索引张量的形状数组地址, 其维度数量与 input_shape_size 相同。

  • core_mask - 核掩码(仅共享存储版本使用)。

输出:
  • output - 输出张量的数据地址, 其形状与 index_shape 保持一致, 数据类型与 input_x 相同。

支持平台:

FT78NE MT7004

备注

  • index 中的取值应满足 0 <= index[...] < input_shape[dim]

  • 输出张量的元素个数等于 index_shape 各维度之积。

  • 该算子不对索引顺序进行任何排序或检查。

共享存储版本:

void fp_gatherd_s(float *input_x, int dim, int *index, float *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void dp_gatherd_s(double *input_x, int dim, int *index, double *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void i8_gatherd_s(int8_t *input_x, int dim, int *index, int8_t *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void i16_gatherd_s(int16_t *input_x, int dim, int *index, int16_t *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void i32_gatherd_s(int32_t *input_x, int dim, int *index, int32_t *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void c64_gatherd_s(float *input_x, int dim, int *index, float *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)
void c128_gatherd_s(double *input_x, int dim, int *index, double *output, int *input_shape, int input_shape_size, int *index_shape, int core_mask)

C调用示例:

 1// FT78NE 示例
 2#include <stdio.h>
 3#include <gatherd.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input_x = (float *)0xA0000000;   // input_x 在 DDR 空间
 7    float *output  = (float *)0xB0000000;
 8    int *index     = (int *)0xA1000000;
 9
10    int input_shape[] = {4, 8, 16};
11    int index_shape[] = {4, 8, 16};
12    int input_shape_size = 3;
13    int dim = 1;
14    int core_mask = 0xff;
15
16    fp_gatherd_s(input_x, dim, index, output, input_shape, input_shape_size, index_shape, core_mask);
17    return 0;
18}

私有存储版本:

void fp_gatherd_p(float *input_x, int dim, int *index, float *output, int *input_shape, int input_shape_size, int *index_shape)
void dp_gatherd_p(double *input_x, int dim, int *index, double *output, int *input_shape, int input_shape_size, int *index_shape)
void i8_gatherd_p(int8_t *input_x, int dim, int *index, int8_t *output, int *input_shape, int input_shape_size, int *index_shape)
void i16_gatherd_p(int16_t *input_x, int dim, int *index, int16_t *output, int *input_shape, int input_shape_size, int *index_shape)
void i32_gatherd_p(int32_t *input_x, int dim, int *index, int32_t *output, int *input_shape, int input_shape_size, int *index_shape)
void c64_gatherd_p(float *input_x, int dim, int *index, float *output, int *input_shape, int input_shape_size, int *index_shape)
void c128_gatherd_p(double *input_x, int dim, int *index, double *output, int *input_shape, int input_shape_size, int *index_shape)

C调用示例:

 1// MT7004 示例
 2#include <stdio.h>
 3#include <gatherd.h>
 4
 5int main(int argc, char* argv[]) {
 6    float *input_x = (float *)0x10000000;   // input_x 在 L2 空间
 7    float *output  = (float *)0x10010000;
 8    int *index     = (int *)0x10020000;
 9
10    int input_shape[] = {2, 4};
11    int index_shape[] = {2, 4};
12    int input_shape_size = 2;
13    int dim = 0;
14
15    fp_gatherd_p(input_x, dim, index, output, input_shape, input_shape_size, index_shape);
16    return 0;
17}